#include <stdlib.h>

#include "png.h"

extern long image_memory(int numbytes, void *data_ptr);
extern void throw_error(char *filename, char *error);

#define PNG_SIG_BYTES 8

static void
premultiply(png_structp png_ptr, png_row_infop info, png_bytep data)
{
	int width = info->width;
	int i;
	
	if(info->channels == 4){
		for(i=0; i<width; i++){
			float alpha = data[i*4 + 3]/255.0;
			data[i*4 + 0] *= alpha;
			data[i*4 + 1] *= alpha;
			data[i*4 + 2] *= alpha;
		}
	} else {
		for(i=0; i<width; i++){
			float alpha = data[i*2 + 1]/255.0;
			data[i*4] *= alpha;
		}
	}
}

long
load_png(char *name, int rgb, int alpha, int premult, int *width, int *height)
{
	FILE *png_file = fopen(name, "rb");
	if(!png_file)
		throw_error(name, "File can't be opened.");
	
	unsigned char header[PNG_SIG_BYTES];
	
	fread(header, 1, PNG_SIG_BYTES, png_file);
	if(png_sig_cmp(header, 0, PNG_SIG_BYTES))
		throw_error(name, "Not a PNG file.");
	
	png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, NULL, NULL, NULL);
	if(!png_ptr)
		throw_error(name, "LibPNG error when loading file.");
	
	png_infop info_ptr = png_create_info_struct(png_ptr);
	if(!info_ptr)
		throw_error(name, "LibPNG error when loading file.");
	
	png_infop end_info = png_create_info_struct(png_ptr);
	if(!end_info)
		throw_error(name, "LibPNG error when loading file.");
	
	if(setjmp(png_jmpbuf(png_ptr)))
		throw_error(name, "LibPNG error when loading file.");
	
	png_init_io(png_ptr, png_file);
	png_set_sig_bytes(png_ptr, PNG_SIG_BYTES);
	png_read_info(png_ptr, info_ptr);
	
	*width = png_get_image_width(png_ptr, info_ptr);
	*height = png_get_image_height(png_ptr, info_ptr);
	
	png_uint_32 bit_depth, color_type;
	bit_depth = png_get_bit_depth(png_ptr, info_ptr);
	color_type = png_get_color_type(png_ptr, info_ptr);
	
	if(color_type == PNG_COLOR_TYPE_GRAY && bit_depth < 8)
		png_set_gray_1_2_4_to_8(png_ptr);
	
	if (bit_depth == 16)
		png_set_strip_16(png_ptr);
	
	if(rgb)
	{
		if(color_type == PNG_COLOR_TYPE_PALETTE)
			png_set_palette_to_rgb(png_ptr);
		else if(color_type == PNG_COLOR_TYPE_GRAY ||
				color_type == PNG_COLOR_TYPE_GRAY_ALPHA)
		{
			png_set_gray_to_rgb(png_ptr);
		}
	}
	else
	{
		if(color_type == PNG_COLOR_TYPE_PALETTE)
			throw_error(name, "Paletted PNG to grayscale conversion not supported.");
		else if(color_type == PNG_COLOR_TYPE_RGB ||
				color_type == PNG_COLOR_TYPE_RGB_ALPHA)
		{
			png_set_rgb_to_gray_fixed(png_ptr, 1, -1, -1);
		}
	}
	
	if (alpha)
	{
		if (png_get_valid(png_ptr, info_ptr, PNG_INFO_tRNS))
			png_set_tRNS_to_alpha(png_ptr);
		else
			png_set_filler(png_ptr, 0xff, PNG_FILLER_AFTER);
	}
	else
	{
		if (color_type & PNG_COLOR_MASK_ALPHA)
			png_set_strip_alpha(png_ptr);
	}
	
	if(premult)
		png_set_read_user_transform_fn(png_ptr, &premultiply);
    
	png_read_update_info(png_ptr, info_ptr);
	
	png_uint_32 rowbytes = png_get_rowbytes(png_ptr, info_ptr);
	png_uint_32 numbytes = rowbytes*(*height);
	
	png_byte *pixels;
	long ruby_value = image_memory(numbytes, &pixels);
	png_byte **row_ptrs = malloc((*height) * sizeof(png_byte*));
	
	int i;
	for (i=0; i<(*height); i++)
		row_ptrs[i] = pixels + ((*height) - 1 - i)*rowbytes;
	
	png_read_image(png_ptr, row_ptrs);
	
	free(row_ptrs);
	png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
	fclose(png_file);
	
	return ruby_value;
}
